"""
A function that plots data and allows the user to click to add points and select points
to remove them.  The function returns the list of points.
"""

from matplotlib.widgets import RectangleSelector
from pylab import subplot, plot, show, ylim, axvline, axhline
import copy

#----------------------------------------------------------------
#-- CLASSES AND FUNCTIONS --
#----------------------------------------------------------------
def rectangle_select_callback(event1, event2):
    'event1 and event2 are the press and release events'
    x1, y1 = int(round(event1.xdata)), event1.ydata
    add_or_remove_pt(x1,y1)
    show()

def add_or_remove_pt(x,y):
    '''
    Maintain a list of points clicked on but not selected out.
    '''
    
    add_or_remove_pt.wave[x] = y
    add_or_remove_pt.line1.set_data(add_or_remove_pt.xlist,\
                                    add_or_remove_pt.wave)
    
        

def wave_maker(wave=None):
    ax=subplot(111)                    # make a new plotingrange

    if not wave==None:
        _wave = copy.deepcopy(wave)
        # plot something... this will be overridden by clicking
        add_or_remove_pt.line1, = plot(  range(len(wave)), wave, linestyle='-', marker='.', linewidth=3, color='red'  )
        ylim(-1,1)
        
    else:
        _wave = [0.0 for i in range(30)]
        add_or_remove_pt.line1, = plot(  range(len(_wave)), _wave, linestyle='-', marker='.', linewidth=3, color='red'  )
        ylim(-1,1)
    # plot vertical and horizontal lines for ease of use.
    for i in xrange(len(_wave)):
        axvline( i,   linewidth=0.5, linestyle=':', color='black')
        axhline( 0.0, linewidth=0.5, linestyle=':', color='black')
    add_or_remove_pt.wave = _wave
    add_or_remove_pt.xlist = range(len(_wave)) 
    # drawtype is 'box' or 'line' or 'none'
    LS = RectangleSelector(ax, rectangle_select_callback,
                      drawtype='box',useblit=True)
    show()
    return add_or_remove_pt.wave
